import numpy as np
import re
import onnx
from pathlib import Path
from config import logger
from collections import defaultdict, OrderedDict

class MindAccMapper:

    @classmethod
    def get_ms_bin_info(cls, ms_bin_file: Path):
        # 从文件名中正则解析出shape和dtype   {op_name}_{input_output_index}_{shape}_{data_type}_{format}.bin
        MS_DATATYPE_TO_NP = {
            "Float64": np.float64,
            "Float32": np.float32,
            "Float16": np.float16,
            "Int64": np.int64,
            "Int32": np.int32,
            "Int8": np.int8,
            "UInt64": np.uint64,
            "UInt8": np.uint8,
            "Bool": np.bool_,
        }
        ms_out_pattern = re.compile(
            r"(?P<op_name>\w+)_(?P<io_flag>(input|output))_(?P<io_index>\d+)_(shape)(_(?P<shape>(\d+(_\d+)*)))?_(?P<data_type>[^_]+)(_(?P<layout>\w+))?"
        )
        shape_str_to_np = lambda shape_str: (tuple(map(int, shape_str.split("_"))) if shape_str else ())
        try:
            node_match = ms_out_pattern.match(ms_bin_file.stem).groupdict()
            node_match["shape"] = shape_str_to_np(node_match["shape"])
            node_match["data_type"] = MS_DATATYPE_TO_NP[node_match["data_type"]]
            return node_match
        except Exception as e:
            logger.warning(f"Can't parse file {ms_bin_file.name}, error: {e}")

    @classmethod
    def read_ms_output(cls, ms_output_path: Path):
        """
        解析结果格式如：
        {1: {'shape': (6,), 'data_type': <class 'numpy.int32'>, 'layout': 'NCHW', 'file': PosixPath('/home/liuhaoyi/oh24/mindacc/output/ssd12/9/    TopK_717_output_1_shape_6_Int32_NCHW.bin')}, 0: {'shape':   (6,), 'data_type': <class 'numpy.float32'>, 'layout': 'NCHW', 'file': PosixPath('/home/liuhaoyi/oh24/   mindacc/output/ssd12/9/TopK_717_output_0_shape_6_Float32_NCHW.bin')}}
        """

        ms_output = defaultdict(dict)
        for file in ms_output_path.iterdir():
            if file.is_file() and file.suffix == ".bin":
                x = MindAccMapper.get_ms_bin_info(file)
                if x["io_flag"] == "output":
                    index = int(x["io_index"])
                    ms_output[x["op_name"]][index] = {
                        "shape": x["shape"],
                        "data_type": x["data_type"],
                        "layout": x["layout"],
                        "file": file,
                    }
        return ms_output

    def __init__(
            self,
            onnx_model: str,
            onnx_dump_file: Path,
            ms_dump_dir: Path,
            extra_rules: dict = {},
    ) -> None:
        self.onnx_model = onnx.load(onnx_model)
        self.onnx_dump = np.load(onnx_dump_file)
        self.ms_dump = MindAccMapper.read_ms_output(Path(ms_dump_dir))
        self.extra_rules = extra_rules
        self.map = {}

    def simple_map(self):
        hwc2chw = lambda shape: (shape[0], shape[3], shape[1], shape[2]) if len(shape) == 4 else shape
        chw2hwc = lambda shape: (shape[0], shape[2], shape[3], shape[1]) if len(shape) == 4 else shape
        dtype_matcher = lambda dtype1, dtype2: (dtype1 == dtype2 or
                                                (dtype1 in [np.int32, np.int64] and dtype2 in [np.int32, np.int64]))
        shape_matcher = lambda shape1, shape2: (shape1 == shape2 or
                                                (len(shape1) == 4 and len(shape2) == 4 and hwc2chw(shape1) == shape2) or
                                                (len(shape1) == 4 and len(shape2) == 4 and chw2hwc(shape1) == shape2))

        for ms_node in self.ms_dump:
            # 在 onnx model 中找到对应的 node
            onnx_node = next((x for x in self.onnx_model.graph.node if x.name == ms_node), None)
            if onnx_node is None:
                logger.warning(f"Can't find node {ms_node} in onnx model")
                continue
            # 为每一个 output 找到对应的 onnx output
            onnx_outputs = onnx_node.output
            matched_onnx = []
            for i, ms_output_info in self.ms_dump[ms_node].items():
                for j, onnx_output in enumerate(onnx_outputs):
                    # 条件1: 形状匹配
                    shape_match = shape_matcher(ms_output_info["shape"], self.onnx_dump[onnx_output].shape)
                    # 条件2: 数据类型匹配
                    dtype_match = dtype_matcher(ms_output_info["data_type"], self.onnx_dump[onnx_output].dtype)
                    # 条件3: ONNX输出未匹配
                    not_matched = onnx_output not in matched_onnx
                    # 满足所有条件
                    match_condition = shape_match and dtype_match and not_matched
                    if match_condition:
                        matched_onnx.append(onnx_output)
                        self.ms_dump[ms_node][i]["onnx_output"] = onnx_output

        self.map = OrderedDict()
        for onnx_node in self.onnx_model.graph.node:
            name = onnx_node.name
            for i, ms_output_info in self.ms_dump[name].items():
                if "onnx_output" in ms_output_info:
                    file_name = ms_output_info['file'].name
                    self.map[file_name] = ms_output_info['onnx_output']
                else:
                    logger.warning(f"Can't find node {ms_node} output {i} in onnx model")
        return self.map

    def get_map_result(self):
        # 获取映射成功率
        maped_count = len(self.map)
        all_count = sum([len(self.ms_dump[node]) for node in self.ms_dump])
        map_rate = maped_count / all_count
        # 获取未映射列表
        unmap_list = []
        for ms_node in self.ms_dump:
            for i, ms_output_info in self.ms_dump[ms_node].items():
                if "onnx_output" not in ms_output_info:
                    file_name = ms_output_info['file'].name
                    unmap_list.append(file_name)
        # 获取映射列表详情
        map_list = {}
        for i in self.map:
            tuple_to_str = lambda t: '_'.join(map(str, t))
            v = f"{self.map[i]}_{tuple_to_str(self.onnx_dump[self.map[i]].shape)}_{self.onnx_dump[self.map[i]].dtype}"
            map_list[i] = v
        return maped_count, all_count, map_rate, unmap_list, map_list
import numpy as np
import re
import onnx
from pathlib import Path
from config import logger
from collections import defaultdict, OrderedDict

class MindAccMapper:

    @classmethod
    def get_ms_bin_info(cls, ms_bin_file: Path):
        # 从文件名中正则解析出shape和dtype   {op_name}_{input_output_index}_{shape}_{data_type}_{format}.bin
        MS_DATATYPE_TO_NP = {
            "Float64": np.float64,
            "Float32": np.float32,
            "Float16": np.float16,
            "Int64": np.int64,
            "Int32": np.int32,
            "Int8": np.int8,
            "UInt64": np.uint64,
            "UInt8": np.uint8,
            "Bool": np.bool_,
        }
        ms_out_pattern = re.compile(
            r"(?P<op_name>\w+)_(?P<io_flag>(input|output))_(?P<io_index>\d+)_(shape)(_(?P<shape>(\d+(_\d+)*)))?_(?P<data_type>[^_]+)(_(?P<layout>\w+))?"
        )
        shape_str_to_np = lambda shape_str: (tuple(map(int, shape_str.split("_"))) if shape_str else ())
        try:
            node_match = ms_out_pattern.match(ms_bin_file.stem).groupdict()
            node_match["shape"] = shape_str_to_np(node_match["shape"])
            node_match["data_type"] = MS_DATATYPE_TO_NP[node_match["data_type"]]
            return node_match
        except Exception as e:
            logger.warning(f"Can't parse file {ms_bin_file.name}, error: {e}")

    @classmethod
    def read_ms_output(cls, ms_output_path: Path):
        """
        解析结果格式如：
        {1: {'shape': (6,), 'data_type': <class 'numpy.int32'>, 'layout': 'NCHW', 'file': PosixPath('/home/liuhaoyi/oh24/mindacc/output/ssd12/9/    TopK_717_output_1_shape_6_Int32_NCHW.bin')}, 0: {'shape':   (6,), 'data_type': <class 'numpy.float32'>, 'layout': 'NCHW', 'file': PosixPath('/home/liuhaoyi/oh24/   mindacc/output/ssd12/9/TopK_717_output_0_shape_6_Float32_NCHW.bin')}}
        """

        ms_output = defaultdict(dict)
        for file in ms_output_path.iterdir():
            if file.is_file() and file.suffix == ".bin":
                x = MindAccMapper.get_ms_bin_info(file)
                if x["io_flag"] == "output":
                    index = int(x["io_index"])
                    ms_output[x["op_name"]][index] = {
                        "shape": x["shape"],
                        "data_type": x["data_type"],
                        "layout": x["layout"],
                        "file": file,
                    }
        return ms_output

    def __init__(
            self,
            onnx_model: str,
            onnx_dump_file: Path,
            ms_dump_dir: Path,
            extra_rules: dict = {},
    ) -> None:
        self.onnx_model = onnx.load(onnx_model)
        self.onnx_dump = np.load(onnx_dump_file)
        self.ms_dump = MindAccMapper.read_ms_output(Path(ms_dump_dir))
        self.extra_rules = extra_rules
        self.map = {}

    def simple_map(self):
        hwc2chw = lambda shape: (shape[0], shape[3], shape[1], shape[2]) if len(shape) == 4 else shape
        chw2hwc = lambda shape: (shape[0], shape[2], shape[3], shape[1]) if len(shape) == 4 else shape
        dtype_matcher = lambda dtype1, dtype2: (dtype1 == dtype2 or
                                                (dtype1 in [np.int32, np.int64] and dtype2 in [np.int32, np.int64]))
        shape_matcher = lambda shape1, shape2: (shape1 == shape2 or
                                                (len(shape1) == 4 and len(shape2) == 4 and hwc2chw(shape1) == shape2) or
                                                (len(shape1) == 4 and len(shape2) == 4 and chw2hwc(shape1) == shape2))

        for ms_node in self.ms_dump:
            # 在 onnx model 中找到对应的 node
            onnx_node = next((x for x in self.onnx_model.graph.node if x.name == ms_node), None)
            if onnx_node is None:
                logger.warning(f"Can't find node {ms_node} in onnx model")
                continue
            # 为每一个 output 找到对应的 onnx output
            onnx_outputs = onnx_node.output
            matched_onnx = []
            for i, ms_output_info in self.ms_dump[ms_node].items():
                for j, onnx_output in enumerate(onnx_outputs):
                    # 条件1: 形状匹配
                    shape_match = shape_matcher(ms_output_info["shape"], self.onnx_dump[onnx_output].shape)
                    # 条件2: 数据类型匹配
                    dtype_match = dtype_matcher(ms_output_info["data_type"], self.onnx_dump[onnx_output].dtype)
                    # 条件3: ONNX输出未匹配
                    not_matched = onnx_output not in matched_onnx
                    # 满足所有条件
                    match_condition = shape_match and dtype_match and not_matched
                    if match_condition:
                        matched_onnx.append(onnx_output)
                        self.ms_dump[ms_node][i]["onnx_output"] = onnx_output

        self.map = OrderedDict()
        for onnx_node in self.onnx_model.graph.node:
            name = onnx_node.name
            for i, ms_output_info in self.ms_dump[name].items():
                if "onnx_output" in ms_output_info:
                    file_name = ms_output_info['file'].name
                    self.map[file_name] = ms_output_info['onnx_output']
                else:
                    logger.warning(f"Can't find node {ms_node} output {i} in onnx model")
        return self.map

    def get_map_result(self):
        # 获取映射成功率
        maped_count = len(self.map)
        all_count = sum([len(self.ms_dump[node]) for node in self.ms_dump])
        map_rate = maped_count / all_count
        # 获取未映射列表
        unmap_list = []
        for ms_node in self.ms_dump:
            for i, ms_output_info in self.ms_dump[ms_node].items():
                if "onnx_output" not in ms_output_info:
                    file_name = ms_output_info['file'].name
                    unmap_list.append(file_name)
        # 获取映射列表详情
        map_list = {}
        for i in self.map:
            tuple_to_str = lambda t: '_'.join(map(str, t))
            v = f"{self.map[i]}_{tuple_to_str(self.onnx_dump[self.map[i]].shape)}_{self.onnx_dump[self.map[i]].dtype}"
            map_list[i] = v
        return maped_count, all_count, map_rate, unmap_list, map_list
